/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "csr_matmul.h"

namespace YHAMG
{

void CSR_MatMul(const CSR_Matrix& A, const CSR_Matrix& B, CSR_Matrix& C)
{
	int n = A.size[0];
	int m = B.size[1];
	int* Ap = A.rowptr;
	int* Ai = A.colind;
	double* Av = A.values;
	int* Bp = B.rowptr;
	int* Bi = B.colind;
	double* Bv = B.values;

	int* Cp = new int[n + 1];
	Cp[0] = 0;
#ifdef USE_OPENMP
#pragma omp parallel
#endif
	{
		int* w = new int[m];
		for (int i = 0; i < m; ++i)
			w[i] = -1;

#ifdef USE_OPENMP
#pragma omp for schedule(guided)
#endif
		for (int i = 0; i < n; ++i)
		{
			int cnt = 0;

			for (int k = Ap[i]; k < Ap[i + 1]; ++k)
			{
				int kcol = Ai[k];

				for (int j = Bp[kcol]; j < Bp[kcol + 1]; ++j)
				{
					int jcol = Bi[j];

					if (w[jcol] != i)
					{
						w[jcol] = i;
						++cnt;
					}
				}
			}

			Cp[i + 1] = cnt;
		}

		delete[] w;
	}

	for (int i = 0; i < n; ++i)
		Cp[i + 1] += Cp[i];

	int* Ci = new int[Cp[n]];
	double* Cv = new double[Cp[n]];

#ifdef USE_OPENMP
#pragma omp parallel
#endif
	{
		int* w = new int[m];
		for (int i = 0; i < m; ++i)
			w[i] = -1;
#ifdef USE_OPENMP
#pragma omp for schedule(guided)
#endif
		for (int i = 0; i < n; ++i)
		{
			for (int k = Ap[i], r = Cp[i], r0 = Cp[i]; k < Ap[i + 1]; ++k)
			{
				int kcol = Ai[k];
				double kval = Av[k];

				for (int j = Bp[kcol]; j < Bp[kcol + 1]; ++j)
				{
					int jcol = Bi[j];
					double jval = Bv[j];

					if (w[jcol] < r0)
					{
						w[jcol] = r;
						Ci[r] = jcol;
						Cv[r++] = kval * jval;
					}
					else
						Cv[w[jcol]] += kval * jval;
				}
			}
		}

		delete[] w;
	}

	C.Free();

	C.size[0] = n;
	C.size[1] = m;
	C.rowptr = Cp;
	C.colind = Ci;
	C.values = Cv;
}

}